import json
import os
from datetime import datetime
from typing import List, Dict, Optional
import logging
import numpy as np

class CustomJSONEncoder(json.JSONEncoder):
    """自定义JSON编码器，处理Numpy数据类型"""
    def default(self, obj):
        if isinstance(obj, (np.integer, np.floating, np.bool_)):
            return obj.item()
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        return super(CustomJSONEncoder, self).default(obj)

class JSONGenerator:
    """JSON输出生成器"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
    
    def generate_flowchart_json(self, 
                               image_path: str,
                               nodes: List[Dict],
                               connections: List[Dict],
                               analysis_result: Dict,
                               metadata: Dict) -> Dict:
        """生成完整的流程图JSON描述"""
        self.logger.info("开始生成JSON输出")
        
        # 构建节点列表
        json_nodes = []
        for node in nodes:
            json_node = {
                'id': node['id'],
                'type': node['type'],
                'text': node.get('text', ''),
                'position': node['position'],
                'confidence': node.get('confidence', 0.0),
                'text_confidence': node.get('text_confidence', 0.0)
            }
            
            # 添加连接关系
            json_node['next'] = self._get_next_nodes(node['id'], connections)
            
            # 如果是决策节点，添加分支信息
            if node['type'] == 'decision':
                json_node['branches'] = self._get_decision_branches(node['id'], connections)
            
            json_nodes.append(json_node)
        
        # 构建连接列表
        json_connections = []
        for connection in connections:
            json_connection = {
                'source': connection['source'],
                'target': connection['target'],
                'type': connection.get('type', 'direct'),
                'points': []
            }
            
            # 转换连接点格式
            if 'points' in connection:
                for point in connection['points']:
                    if isinstance(point, tuple):
                        json_connection['points'].append({'x': point[0], 'y': point[1]})
                    else:
                        json_connection['points'].append(point)
            
            # 添加箭头信息
            if 'arrow' in connection:
                json_connection['arrow'] = {
                    'x': connection['arrow'][0],
                    'y': connection['arrow'][1]
                }
            
            json_connections.append(json_connection)
        
        # 构建元数据
        json_metadata = {
            'filename': os.path.basename(image_path),
            'width': metadata.get('processed_size', (0, 0))[1],
            'height': metadata.get('processed_size', (0, 0))[0],
            'creation_time': datetime.now().isoformat() + 'Z',
            'confidence_score': self._calculate_overall_confidence(nodes, connections),
            'processing_info': {
                'original_size': metadata.get('original_size', (0, 0)),
                'skew_angle': metadata.get('skew_angle', 0.0),
                'quality_score': metadata.get('quality_score', 0.0)
            }
        }
        
        # 构建结构信息
        json_structure = {
            'type': analysis_result.get('type', 'sequence'),
            'subflows': analysis_result.get('subflows', []),
            'cycles': analysis_result.get('cycles', []),
            'branches': analysis_result.get('branches', []),
            'summary': self._generate_summary(analysis_result)
        }
        
        # 组装完整的JSON
        result = {
            'metadata': json_metadata,
            'nodes': json_nodes,
            'connections': json_connections,
            'structure': json_structure
        }
        
        self.logger.info("JSON生成完成")
        return result
    
    def _get_next_nodes(self, node_id: str, connections: List[Dict]) -> List[str]:
        """获取节点的下一个节点列表"""
        next_nodes = []
        for connection in connections:
            if connection['source'] == node_id:
                next_nodes.append(connection['target'])
        return next_nodes
    
    def _get_decision_branches(self, node_id: str, connections: List[Dict]) -> Dict:
        """获取决策节点的分支信息"""
        branches = {}
        
        for connection in connections:
            if connection['source'] == node_id:
                # 根据连接类型确定分支
                if connection.get('type') == 'true_branch':
                    branches['true'] = [connection['target']]
                elif connection.get('type') == 'false_branch':
                    branches['false'] = [connection['target']]
                else:
                    # 默认分支
                    if 'default' not in branches:
                        branches['default'] = []
                    branches['default'].append(connection['target'])
        
        return branches
    
    def _calculate_overall_confidence(self, nodes: List[Dict], connections: List[Dict]) -> float:
        """计算整体置信度"""
        if not nodes:
            return 0.0
        
        # 计算节点检测置信度
        node_confidences = [node.get('confidence', 0.0) for node in nodes]
        avg_node_confidence = sum(node_confidences) / len(node_confidences)
        
        # 计算文本识别置信度
        text_confidences = [node.get('text_confidence', 0.0) for node in nodes if node.get('text')]
        avg_text_confidence = sum(text_confidences) / len(text_confidences) if text_confidences else 0.0
        
        # 计算连接置信度（基于连接数量）
        connection_confidence = min(1.0, len(connections) / max(len(nodes) - 1, 1))
        
        # 综合置信度
        overall_confidence = (avg_node_confidence * 0.4 + 
                            avg_text_confidence * 0.4 + 
                            connection_confidence * 0.2)
        
        return round(overall_confidence, 2)
    
    def _generate_summary(self, analysis_result: Dict) -> Dict:
        """生成流程摘要"""
        graph_info = analysis_result.get('graph_info', {})
        structure = analysis_result.get('structure', {})
        
        summary = {
            'total_nodes': graph_info.get('total_nodes', 0),
            'total_connections': graph_info.get('total_connections', 0),
            'has_cycles': graph_info.get('has_cycles', False),
            'decision_count': len(structure.get('decision_nodes', [])),
            'start_nodes': structure.get('start_nodes', []),
            'end_nodes': structure.get('end_nodes', []),
            'complexity': self._determine_complexity(graph_info, structure)
        }
        
        return summary
    
    def _determine_complexity(self, graph_info: Dict, structure: Dict) -> str:
        """确定流程复杂度"""
        total_nodes = graph_info.get('total_nodes', 0)
        decision_count = len(structure.get('decision_nodes', []))
        has_cycles = graph_info.get('has_cycles', False)
        
        if total_nodes <= 5 and decision_count <= 1 and not has_cycles:
            return 'simple'
        elif total_nodes <= 15 and decision_count <= 3:
            return 'medium'
        else:
            return 'complex'
    
    def save_json_to_file(self, json_data: Dict, output_path: str) -> bool:
        """保存JSON到文件"""
        try:
            # 确保输出目录存在
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # 保存JSON文件
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(json_data, f, ensure_ascii=False, indent=2, cls=CustomJSONEncoder)
            
            self.logger.info(f"JSON文件已保存到: {output_path}")
            return True
            
        except Exception as e:
            self.logger.error(f"保存JSON文件失败: {e}")
            return False
    
    def generate_error_json(self, error_message: str, details: Optional[Dict] = None) -> Dict:
        """生成错误JSON"""
        error_json = {
            'error': True,
            'message': error_message,
            'timestamp': datetime.now().isoformat() + 'Z'
        }
        
        if details:
            error_json['details'] = details
        
        return error_json
    
    def generate_output_path(self, input_path: str, output_dir: Optional[str] = None) -> str:
        """生成输出文件路径"""
        if output_dir is None:
            output_dir = os.path.dirname(input_path)
        
        # 获取输入文件名（不含扩展名）
        base_name = os.path.splitext(os.path.basename(input_path))[0]
        
        # 生成输出文件名
        output_filename = f"{base_name}_parsed.json"
        
        return os.path.join(output_dir, output_filename)
    
    def validate_json_structure(self, json_data: Dict) -> bool:
        """验证JSON结构是否符合规范"""
        required_fields = ['metadata', 'nodes', 'connections', 'structure']
        
        # 检查必需字段
        for field in required_fields:
            if field not in json_data:
                self.logger.error(f"缺少必需字段: {field}")
                return False
        
        # 检查元数据字段
        metadata = json_data['metadata']
        metadata_fields = ['filename', 'width', 'height', 'creation_time', 'confidence_score']
        for field in metadata_fields:
            if field not in metadata:
                self.logger.error(f"元数据缺少字段: {field}")
                return False
        
        # 检查节点结构
        for node in json_data['nodes']:
            node_fields = ['id', 'type', 'text', 'position']
            for field in node_fields:
                if field not in node:
                    self.logger.error(f"节点缺少字段: {field}")
                    return False
        
        # 检查连接结构
        for connection in json_data['connections']:
            connection_fields = ['source', 'target', 'type']
            for field in connection_fields:
                if field not in connection:
                    self.logger.error(f"连接缺少字段: {field}")
                    return False
        
        self.logger.info("JSON结构验证通过")
        return True 